// SPDX-License-Identifier: Mulan PSL v2
/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This software is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *         http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

use std::{
    sync::atomic::{AtomicUsize, Ordering},
    thread,
    time::Duration,
};

use tracing::debug;

use crate::{agent, fault_guard::journal};

static MINIBATCH_COUNT: AtomicUsize = AtomicUsize::new(0);

fn on_minibatch_begin() -> anyhow::Result<()> {
    agent::set_recovery_flag(true);

    while agent::get_active_request_count() > 0 {
        thread::sleep(Duration::from_millis(1));
    }

    MINIBATCH_COUNT.fetch_add(1, Ordering::Relaxed);
    let minibatch_count = MINIBATCH_COUNT.load(Ordering::Relaxed);
    debug!("[Hook] Minibatch {} begins", minibatch_count);

    let old_queue = journal::reset_minibatch_requests()?;
    thread::spawn(move || {
        while let Some(req_bytes) = old_queue.pop() {
            journal::dump_request(&req_bytes).expect("Failed to log minibatch request");
        }
    });

    agent::set_recovery_flag(false);

    Ok(())
}

#[unsafe(no_mangle)]
pub extern "C" fn minibatch_begin_hook() {
    on_minibatch_begin().expect("Failed to execute on_minibatch_begin");
}
